import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
import zipfile
import io
from insightface.app import FaceAnalysis
import cv2
import os
from typing import List, Dict
import uuid

from pymilvus import MilvusClient

from FaceInformationSingUp import sign_up

app = FastAPI()
app2 = FaceAnalysis(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])  # 优先GPU加速
# 允许的图片文件扩展名
ALLOWED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"}

cli = MilvusClient(
    uri='http://39.104.78.210:19530',
    token='root:Milvus',
) # 建立连接

def is_image_filename(filename: str) -> bool:
    """检查文件是否为图片类型"""
    return any(filename.lower().endswith(ext) for ext in ALLOWED_IMAGE_EXTENSIONS)


@app.post("/extract-zip-images")
async def extract_zip_images(zip_file: UploadFile = File(...)):
    # 检查文件类型是否为ZIP
    if not zip_file.filename.endswith('.zip'):
        raise HTTPException(
            status_code=400,
            detail="请上传ZIP文件"
        )

    try:
        # 读取ZIP文件内容
        contents = await zip_file.read()

        # 使用zipfile处理内存中的ZIP文件
        with zipfile.ZipFile(io.BytesIO(contents), 'r') as zip_ref:
            # 获取ZIP文件中所有文件列表
            file_list = zip_ref.namelist()

            # 筛选出图片文件
            image_files = [f for f in file_list if is_image_filename(f)]

            # 如果没有找到图片文件
            if not image_files:
                return {
                    "message": "ZIP文件中未找到图片文件",
                }

            extracted_images = []
            extracted_name = []
            for image_filename in image_files:
                try:
                    # 从ZIP提取文件内容
                    with zip_ref.open(image_filename) as image_file:
                        np_array = np.frombuffer(image_file.read(), np.uint8)
                        image_data = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
                        extracted_name.append(image_filename.split('.')[0])
                        extracted_images.append(image_data)
                except Exception as e:
                    # 记录错误但继续处理其他文件
                    raise HTTPException(
                        status_code=400,
                        detail="无效的ZIP文件"
                    )

            return {
                "message": sign_up(cli, app2, extracted_images, extracted_name, "faceD"),
                "zip_file": zip_file.filename,
            }

    except zipfile.BadZipFile:
        raise HTTPException(
            status_code=400,
            detail="无效的ZIP文件"
        )
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=500,
            detail=f"处理ZIP文件时出错: {str(e)}"
        )


if __name__ == "__main__":
    import uvicorn

    # uvicorn.run(app, host="0.0.0.0", port=8000)